package com.tianzhi.modeldesigner;

import java.util.Collections;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.testcontainers.containers.MongoDBContainer;
import org.testcontainers.utility.DockerImageName;

public class MongoDbTestContainerExtension implements BeforeAllCallback {

    private static final ThreadLocal<MongoDBContainer> THREAD_CONTAINER = new ThreadLocal<>();
    private final long memoryInBytes = Math.round(1024 * 1024 * 1024 * 0.6);
    private final long memorySwapInBytes = Math.round(1024 * 1024 * 1024 * 0.8);
    private final long nanoCpu = Math.round(1_000_000_000L * 0.1);

    @Override
    public void beforeAll(ExtensionContext context) {
        if (null == THREAD_CONTAINER.get()) {
            MongoDBContainer mongoDBContainer = new MongoDBContainer(DockerImageName.parse("mongo:4.4.12"))
                .withReuse(true)
                .withTmpFs(Collections.singletonMap("/testtmpfs", "rw"))
                .withCommand(
                    "--nojournal --wiredTigerCacheSizeGB 0.25 --wiredTigerCollectionBlockCompressor none --slowOpSampleRate 0 --setParameter ttlMonitorEnabled=false --setParameter diagnosticDataCollectionEnabled=false --setParameter logicalSessionRefreshMillis=6000000 --setParameter enableFlowControl=false --setParameter oplogFetcherUsesExhaust=false --setParameter disableResumableRangeDeleter=true --setParameter enableShardedIndexConsistencyCheck=false --setParameter enableFinerGrainedCatalogCacheRefresh=false --setParameter readHedgingMode=off --setParameter loadRoutingTableOnStartup=false --setParameter rangeDeleterBatchDelayMS=2000000 --setParameter skipShardingConfigurationChecks=true --setParameter syncdelay=3600"
                )
                .withCreateContainerCmdModifier(cmd ->
                    cmd.getHostConfig().withMemory(memoryInBytes).withMemorySwap(memorySwapInBytes).withNanoCPUs(nanoCpu)
                );
            mongoDBContainer.start();
            THREAD_CONTAINER.set(mongoDBContainer);
        }
    }

    public static ThreadLocal<?> getThreadContainer() {
        return THREAD_CONTAINER;
    }
}
